#/usr/bin/python
# -*- coding: utf-8 -*-
'''

This module supplies finite element discretization of a geometric domain.

Created on 5 mars 2012

@author: Ludovic Barrière and Steven Marguet, Institut Clément Ader, 135 av. de Rangueil, 31077 Toulouse, France

© copyright:
    
    TODO
'''
import sys
sys.path.append('../../icafe')
import parameters

from mpi4py import MPI
import scipy
import scipy.sparse as sp
import numpy
import numpy as np
import scipy.linalg as la
import copy

from ioparser.parserModel.parserPatran import Structural_Part
import cPickle as pickler
from elements import *

FLOAT = "float64"

eltype_to_modelclass = {
    'S3': Tri3(),
    'S4': Quad4(),
    'S6': Tri6(),
    'S8': Quad8(),
    'B2': Beam2(),
    'B3': Beam3(),
    'R2': Rod2(),          
}

class DefaultMesh(object):
    '''
    A Mesh defines a discrete counterpart of a geometric domain with physical properties. 
    Equilibrium equation (Principle of Virtual Work) is discretized also on this mesh.
    In a programming point of view, mesh if made of "element" entities and "node" entities. 
    The left hand side (Lhs) of the equilibrium equation is computed by an assembly of different
    element contributions in a square matrix (symmetric definite positive in the case of classical finite 
    element method). The right hand side (Rhs) of the equilibrium equation is computed in the same way.
    This Mesh class is designed in order to contain data of any 3D mesh, made of any element type (see the "Element" class below).
    Enriched interpolation can be set easily. 
    
    '''
    def __init__(self, UId, **kwargs):
        '''
            @param UId: str or int.
        '''
        # mesh prop:
        self.UId = UId
        self.ndofs_per_node = 6
        
        # Node entities:
        self.nodesCoo = []
        self.nodesElements = []
        # Mapping of nodes:
        self.nodesMap = {"U2C": {}, "C2U": [],
                         "COO2C": {},"C2COO": []}
        self.nnodes = 0
        # Element entities:
        self.elementsNodes = []
        self.elementsAO = []
        self.elementsFrame = []
        self.elementsProp = []
        self.elementsType = []
        # Mapping of elements:
        self.elementsMap = {"U2C": {}, "C2U": []}
        self.nElem = 0
        
        # Property entities:
        self.materials = {}
        self.sections = {}
        self.properties = [] # list of couple of section and material id
        # mapping of properties:
        self.propertiesMap = {"U2C": {}, "C2U": []}
        
        # Enrichment entities:
        self.enrichments = []
        
        # Mapping of enrichment:
        self.enrichmentsMap = {
                             "U2C": {},"C2U": [], # mapping user name of enrichment <-> code id of enrichment
                             "C2DOF":[],"DOF2C":{}, # mapping code id of enrichment <-> code id of 
                             "C2ELT":[],"ELT2C":{},
                             "C2DOFTYPE":[],"DOFTYPE2C":{0:[],1:[],2:[],3:[],4:[],5:[]},
                             }
        # Subspace entities:
        self.subspaces = []
        
        # Mapping of subspace:
        self.subspacesMap = {"U2C": {}, "C2U": []}

        # degrees of freedom:
        self.dofs = [] #sp.lil_matrix()
        self.duals = []
        
        self.constraints=None
        
        # fields history:
        self.fields = Fields(self.UId)
        # quantities history:
        self.quantities = Quantities(self.UId)
        self.error = 0
#        self.history_dofs = []
#        self.history_dofs_Map = {'U2C':{}, 'C2U':[]}
        
        # Operators:
        self.lhs = {
                    "M":None,
                    "K":None, 
                    "K_linear":None,
                    "K_nonlinear":None,
                    "Kt":None, 
                    "Kt_linear":None, 
                    "Kt_geometric":None, 
                    "Kt_nonlinear":None
                    }
        self.rhs = {
                    "Fext":None, 
                    "R":None
                    }
        
        # hp-levels:
        self.hLevel = 0
        self.pLevel = 0
    
    
        args = self.__dict__ 
        
        for k in kwargs.keys():
            if k in args.keys():
#                print 'setting arg ', k, ' with ', kwargs[k]
                self.__setattr__(k,kwargs[k])
    
    def __str__(self):
        
        r = '\n' + 'Mesh ' + str(self.UId) + ': \n'
        r += '  - ' + str(self.nnodes) + ' nodes' + '\n'
        r += '  - ' + str(self.ndofs_per_node) + ' d.o.f./node'  #+ '\n'
        return r
    
    def initConstraints(self):
        ''' 
            Initialize constraint operators
        '''
        self.constraints=None
    
    def setConstraints(self, loadcases, bcs, lcs):
        ''' 
            Set constraints
        '''
        self.bC = bcs
        self.lC = lcs
        
        nnodes, ndofs_per_node = self.getShapeDofs()
        ndofs = nnodes*ndofs_per_node
        
        # Initialize boundary conditions
        Cd = []
        d  = []
        # Initialize loads
        Cf = []
        f  = []
        
        for loadcase in loadcases:
            
            # Boundary conditions
            d.append(None)
            Cd.append(None)
            
            for bcId in loadcase.bcIds:
                try:
                    bc = bcs[bcId]
                except KeyError:
                    if self.log:
                        self.log.info('Boundary condition '+str(bcId)+' is not in domain '+ str(self.UId))
                    continue
                
                
                ## WARNING: FOLLOWING IS RESTRICTED TO SIMPLE BOUNDARY CONDITIONS
                ## DEFINED BY NODES AND SIMPLE CONSTRAINTS ON DOFS.
                
                # get nodes:
#                print "edge of bc: ", bc.edge.UId
                cnodes = self.getNodesFromEdge(edge = bc.edge) 
                ncnodes = len(cnodes)
#                print cnodes
                
                if not ncnodes: # Boundary condition does not concern the domain
                    # remove bc from domain.bC:
                    del bcs[bcId]
                    if self.log:
                        self.log.info('Boundary condition '+str(bcId)+' has been removed from domain '+str(self.UId)) 
                    continue
                # add constraints on dofs:
                
                
                for dof, val in bc.constraints.items():
                    if val is not None: 
                        
                        # for each constrained node, adds line to contraint matrix:
                        for cnd in cnodes:
                            c = sp.lil_matrix((1,ndofs))
                            coeff = [1]
                            index = [cnd*ndofs_per_node + dof]
                            # gets enrichment functions of dof:
                            efuncIds = self.enrichmentsMap["DOFTYPE2C"][dof]
                            
                            for fid in efuncIds:
                                # gets values of enrichment function at node
                                func = self.enrichments[fid]
                                coeff.append(func(self.nodesCoo[cnd]))
                                # adds index of enriched dof:
                                pdof = self.enrichmentsMap["C2DOF"][fid]
                                index.append(cnd*ndofs_per_node + pdof)
                
                        
                            # set value of coefficient in extraction matrix:
                            c[0,index] = sp.lil_matrix(coeff)
                            try:
                                d[-1] = sp.bmat([[d[-1]],[val]])
                                Cd[-1] = sp.bmat([[Cd[-1]],[c]])
                            except:
                                d[-1] = sp.lil_matrix([[val]])
                                Cd[-1]= c
                                
                                
            # Loads
            for loadId in loadcase.loadIds:
                # gets load:
                try:
                    load = lcs[loadId]
                except KeyError:
                    if self.log:
                        self.log.info('Load condition '+str(loadId)+' is not in domain '+ str(self.UId))
                    continue
                # gets nodes:
#                print "edge of load: ", load.edge.UId
                cnodes = self.getNodesFromEdge(edge = load.edge) 
                
                ncnodes = len(cnodes)
                
                if not ncnodes: # Load condition does not concern the domain
                    # remove bc from domain.bC:
                    del lcs[loadId]
                    if self.log:
                        self.log.info('Load condition '+str(loadId)+' has been removed from domain '+str(self.UId)) 
                    continue
                
                # init extractor:
                Cf.append(sp.lil_matrix((1,ndofs)))
                f.append(0)
            
                
                for dof, val in load.loads.items():
                    if val is not None:
                        
                        for cnd in cnodes:
                            c = sp.lil_matrix((1,ndofs))
                            
                            c[0,cnd*ndofs_per_node+dof]=1
                            
                            if Cf[-1].getnnz()!=0:
                                f[-1] = sp.bmat([[f[-1]],[val]])
                                Cf[-1] = sp.bmat([[Cf[-1]],[c]])
                            else:
                                f[-1] = sp.lil_matrix([[val]])
                                Cf[-1] = c
        
        constraints = Constraints(Cd, d, Cf, f, self.UId)
        
        self.constraints=constraints
        
        return constraints
    
    def getNodesFromEdge(self, edge):
        '''
            Return a list of nodes belonging to the edge
            @param edge: Instance of class Edge.
        '''
        # init list of nodes:
        nodes = []
        
        try:
            nodes = edge.getNodesFromMesh(self) 
        except KeyError: # if edge does not correspond to an edge of the domain:
            nodes = []
        
        return nodes
           
    # Add/Remove methods:
    def addNode(self, UId, coord):
        '''
            @param UId: str or int.
        '''
        from numpy import around
        param = parameters.load()
        dec = param.DECIMALS
        
        cid = len(self.nodesMap['C2U'])
        
        try:
            cid = self.nodesMap["U2C"][UId]
            return 0
        except KeyError:
            pass
        # Mapping:
        self.nodesMap["U2C"][UId]=cid
        self.nodesMap["C2U"].append(UId)
        ncoo = tuple(around(coord,decimals=dec).tolist())
        self.nodesMap["COO2C"][ncoo]=cid
        self.nodesMap["C2COO"].append(ncoo)
        # Coordinates:
        self.nodesCoo.append(coord)
        # cross referecing:
        self.nodesElements.append([]) # List of User Id of elements, to which the node belongs.
        self.nnodes += 1
    
    def addElement(self, UId, nodes, type, propId=None):
        '''
            @param UId: str or int.
            @param nodes: list of nodes User Ids
        '''
        cid = len(self.elementsMap['C2U'])
        try: 
            cid = self.elementsMap["U2C"][UId]
            return 0
        except KeyError:
            pass
        # mapping
        self.elementsMap["U2C"][UId]=cid
        self.elementsMap["C2U"].append(UId)
        # attributes:
        self.elementsType.append(type)
        self.elementsProp.append(propId)
        self.elementsNodes.append(nodes)   # list of User Id of nodes
        
        # cross referencing:
        for nid in nodes:
            try:
                nid = self.nodesMap['U2C'][nid]
                self.nodesElements[nid].append(UId)
            except KeyError:
                nid = self.elementsNodes[-1].index(nid)
                self.elementsNodes[-1][nid]=None
                #raise KeyError('Problem with nodes:' + repr(self.nodesMap['U2C'])+  repr(nodes))
        self.nElem += 1   
    
    def rmNode(self, UId):
        '''
            @param UId: str or int.
        '''
        pass
    
    def rmElement(self, UId):
        '''
            @param UId: str or int.
        '''
        pass
    
    def addMaterial(self, UId, mat):
        '''
            @param UId: str or int.
            @param mat: Material
        '''
        self.materials[UId] = mat        
        
    def addProperty(self, UId, prop):
        '''
            @param UId: str of int.
            @param prop: Property.
        '''
        cid = len(self.propertiesMap['C2U'])
        
        self.properties.append(prop)
        self.propertiesMap['U2C'][UId]=cid
        self.propertiesMap['C2U'].append(UId)
    
    def rmProperty(self, UId):
        '''
            @param UId: str or int.
        '''
        pass
    
    def addDof(self, enrichmentCId=None):
        '''
        
        '''
        shape = self.dofs.get_shape()
        
        new_dofs = sp.lil_matrix((shape[0],1),dtype=parameters.FLOAT)
        
        if enrichmentCId is not None:
            self.enrichmentsMap["C2DOF"][enrichmentCId].append(shape[1])
            self.enrichmentsMap["DOF2C"][shape[1]]=enrichmentCId

        
        self.dofs = sp.hstack([self.dofs,new_dofs])
        self.ndofs_per_node += 1
        
        return self.getDofs()
            
    def rmDof(self, dof):
        '''
        
        '''
        
        
        self.ndofs_per_node -= 1
        
        return self.getDofs()
    
    def addEnrichment(self, function, dofs='012345'):
        '''
        
        '''
        eid = len(self.enrichments)
        self.enrichments.append(function)
        
        self.enrichmentsMap["C2DOFTYPE"].append([])
        
        for dof in dofs:
            self.addDof()
            self.enrichmentsMap["C2DOFTYPE"][eid]+=int(dof)
            self.enrichmentsMap["DOFTYPE2C"][int(dof)].append(eid)
        
        
    def rmEnrichment(self):
        '''
        
        '''
        pass
        
    # Get/Set methods:
    def getNodes(self, nodelistCId=None, nodelistUId=None, elemCId=None, elemUId=None, mode='user'):
        '''
            
            @param mode: 'user' or 'code', defines wether Ids of nodes/Elements output is UId of CId
        '''
        nodes = {"UId":[],"coord":[],"dofs":[],"elements":[]}
        all = True
        
        if nodelistCId:
            all = False
            for nid in nodelistCId:
                nodes["UId"].append(self.nodesMap['C2U'][nid])
                nodes["coord"].append(self.nodesCoo[nid])
                nodes["elements"].append(self.nodesElements[nid])
                
                    
        if nodelistUId:
            all = False
            for nid in nodelistUId:
                nodes["UId"].append(nid)
                nodes["coord"].append(self.nodesCoo[self.nodesMap['U2C'][nid]])
                nodes["elements"].append(self.nodesElements[self.nodesMap['U2C'][nid]])
        
        if elemCId:
            all = False
            for eid in elemCId:
                for nid in self.elementsNodes[eid]:
                    nodes["UId"].append(self.nodesMap['C2U'][nid])
                    nodes["coord"].append(self.nodesCoo[nid])
                    nodes["elements"].append(self.nodesElements[nid])
            
        if elemUId:
            all = False
            for eid in elemUId:
                for nid in self.elementsNodes[self.elementsMap['U2C'][eid]]:
                    nodes["UId"].append(self.nodesMap['C2U'][nid])
                    nodes["coord"].append(self.nodesCoo[nid])
                    nodes["elements"].append(self.nodesElements[nid])
        
        if all:
            for nid in self.nodesMap['C2U']:
                nodes["UId"].append(nid)
                nodes["coord"].append(self.nodesCoo[self.nodesMap['U2C'][nid]])
                nodes["elements"].append(self.nodesElements[self.nodesMap['U2C'][nid]])
                 
                
        return nodes
    
    def getElements(self, elemCId=None, elemUId=None, mode='user'):
        '''
            @param mode: 'user' or 'code', defines wether Ids of nodes/Elements output is UId of CId
        '''
        elements = {"UId":[], "nodes":[], "types":[],"props":[]}
        all = True
        
        if elemCId:
            all = False
            for eid in elemCId:
                elements['UId'].append(self.elementsMap['C2U'][eid])
                if mode == 'user':
                    elements['nodes'].append(self.elementsNodes[eid])
                else:
                    elements['nodes'].append([self.nodesMap['U2C'][nd] for nd in self.elementsNodes[eid]])
                elements['types'].append(self.elementsType[eid])
                elements['props'].append(self.propertiesMap['C2U'][self.elementsProp[eid]])
           
        if elemUId:
            all = False
            for eid in elemUId:
                eid = self.elementsMap['U2C'][eid]
                elements['UId'].append(self.elementsMap['C2U'][eid])
                if mode == 'user':
                    elements['nodes'].append(self.elementsNodes[eid])
                else:
                    elements['nodes'].append([self.nodesMap['U2C'][nd] for nd in self.elementsNodes[eid]])
                elements['types'].append(self.elementsType[eid])
                elements['props'].append(self.propertiesMap['C2U'][self.elementsProp[eid]])
        
        if all:
            for eid in self.elementsMap['C2U']:
                eid = self.elementsMap['U2C'][eid]
                elements['UId'].append(self.elementsMap['C2U'][eid])
                if mode == 'user':
                    elements['nodes'].append(self.elementsNodes[eid])
                else:
                    elements['nodes'].append([self.nodesMap['U2C'][nd] for nd in self.elementsNodes[eid]])
                elements['types'].append(self.elementsType[eid])
                elements['props'].append(self.propertiesMap['C2U'][self.elementsProp[eid]])
            
        return elements    
    
    def getProperties(self, UId=None):
        
        props = {"UId":[],"types":[], "materials":[],"section":[], "S-args":[]}
        all = True
        
        if UId:
            all = False
            props["UId"].append(UId)
            prop = self.properties[self.propertiesMap["U2C"][UId]]
            props["types"].append(prop.type)
            props["materials"].append(prop.material.mid)
            props["S-args"].append(prop.section.getArgs())

        if all:
            for prop in self.properties:
                props["UId"].append(prop.UId)
                props["types"].append(prop.type)
                props["materials"].append(prop.material)
                props["S-args"].append(prop.section.getArgs())
            
        return props
    
    def getMaterials(self, UId=None):
        
        materials = {}
        
        if UId:
            materials = {UId: self.materials[UId]}
        else:
            materials = self.materials
                
        return materials
    
    def getNodesCooAtState(self, state=None, format='output', components='012', dtype='float32'):
        
        if format=='output':
            x = []
            y = []
            z = []
            for nd in self.nodesCoo:
                x.append(nd[0])
                y.append(nd[1])
                z.append(nd[2])
            x = np.array(x, dtype)
            y = np.array(y, dtype)
            z = np.array(z, dtype)
            if state:
                dofs = self.fields.getField('dof', state)
                if '0' in components:
                    x = x + dofs[:,0].transpose().toarray()[0,:]
                if '1' in components:
                    y += dofs[:,1].transpose().toarray()[0,:]
                if '2' in components:
                    z += dofs[:,2].transpose().toarray()[0,:]
            return x, y, z    
        else:
            if not state:
                return self.nodesCoo
            else:
                pass
    
    def getNodesDofs(self, nodes, dofs='012345', complement=False):
        '''
           @return: degrees of freedom of particular nodes.
           @param nodes: may be a list of nodes UId, or nodes coordinates. 
           @param dofs: a string of degrees of freedom to be retained.
        '''
        if isinstance(nodes,str):
            param = parameters.load()
            nodes = param.getPointsOfInterest(nodes)
            
        values = []    
        for nd in nodes:
            try:
                if isinstance(nd,tuple):
                    nd = self.nodesMap['COO2C'][nd]
                else:
                    nd = self.nodesMap['U2C'][nd]
            except: # no such node in mesh
                continue  
            values.append([self.dofs[nd,eval(i)] for i in dofs])    
            
        if complement:
            return values, nodes    
        else:
            return values
    
    def getNodesForce(self, nodes, duals='012345', complement=False):
        '''
           @return: forces at particular nodes.
           @param nodes: may be a list of nodes UId, or nodes coordinates. 
           @param duals: a string of degrees of freedom to be retained.
        '''
        if isinstance(nodes,str):
            param = parameters.load()
            nodes = param.getPointsOfInterest(nodes)
            
        values = []    
        for nd in nodes:
            try:
                if isinstance(nd,tuple):
                    nd = self.nodesMap['COO2C'][str(nd)]
                else:
                    nd = self.nodesMap['U2C'][nd]
            except: # no such node in this mesh, maybe in another one...
                continue  
                
            values.append([self.duals[nd,int(i)] for i in duals])    
        
        if complement:
            return values, nodes    
        else:
            return values
        
    def getError(self, complement=False):
        
        if complements:
            return self.error, None
        else:
            return self.error
        
    def getNodesArray(self, output_type=sp.lil_matrix):
        '''
           Returns nodes coordinates in an array of shape (number of nodes, 3)
        '''
        return output_type(self.nodesCoo)
    
    def getDofs(self, format=True, components='012345'):
        '''
            @param format: True means that the vector of d.o.f will be formatted as a column vector.
            @param components: string of retained components
        '''
        comp = [int(i) for i in components]
        dofs = copy.deepcopy(self.dofs[:,comp])
            
        if format:
            shape = dofs.get_shape()
            
            return dofs.reshape((shape[0]*shape[1],1))    
        else:
            
            return dofs
    
    def getForces(self, format=True, components='012345'):
        '''
            @param format: True means that the vector of duals will be formatted as a column vector.
            @param components: string of retained components
        '''
        comp = [int(i) for i in components]
        duals = copy.deepcopy(self.duals[:,comp])
            
        if format:
            shape = duals.get_shape()
            
            return duals.reshape((shape[0]*shape[1],1))
        else:
            
            return duals
        
    def getCurrentField(self, field, components='012345'):
        
        gettingFields = {
        'dof':['getDofs',False],
        'force':['getForces',False],
        }
        meth = gettingFields[field][0]
        
        if not gettingFields[field][1] is None: 
            return self.__getattribute__(meth)(gettingFields[field][1],components)
        else:
            return self.__getattribute__(meth)(components)
    
    def getCurrentQuantity(self, qoi, args):
        
        gettingQOI = {
        'error':'getError',
        'dof':'getNodesDofs',   
        'force' :'getNodesForces',          
        }
        
        meth = gettingQOI[qoi]
        if args:
            return self.__getattribute__(meth)(*args,complement=True)
        else:
            return self.__getattribute__(meth)(complement=True)
        
    def getShapeDofs(self):
        
        shape = self.dofs.get_shape()
        
        return shape[0],shape[1]   
    
    def getConstraints(self):
        '''
           Returns displacement and load contraints. 
        '''
        return self.constraints
        
    def getRhs(self, rhs):
        '''
        
        '''
        return self.rhs[rhs] 

    def getLhs(self, lhs):
        '''
        
        '''
        return self.lhs[lhs]
    
    def setElementsAO(self):
        '''
        
        '''
        n_elem = len(self.elementsMap["C2U"])
        n_nodes = len(self.nodesMap["C2U"])
        
        self.elementsAO = [None] * n_elem
        
        for e,el_nodes in enumerate(self.elementsNodes):
            
            nn = len(el_nodes)
            
            self.elementsAO[e] = sp.lil_matrix((6*n_nodes,6*nn),dtype = FLOAT)
            
            I = sp.csr_matrix(np.eye(6),dtype = FLOAT).tolil()
            
            for i,n in enumerate(el_nodes):
                
                self.elementsAO[e][6*n:6*n+6,6*i:6*i+6] = I
    # Other methods:           
    def buildDofs(self):
        """
           
        """
        self.nnodes = len(self.nodesMap['C2U'])
        self.dofs = sp.lil_matrix((self.nnodes,self.ndofs_per_node),dtype=parameters.FLOAT)
        self.duals = sp.lil_matrix((self.nnodes,self.ndofs_per_node),dtype=parameters.FLOAT)
    
    
    def upDateDofs(self, ddofs, dduals=None, op='add'):
        '''
        
        '''
        shape = self.dofs.get_shape()
        
        ddofs = ddofs.reshape(shape)
        
        if dduals is not None:
            shape = self.duals.get_shape()
            dduals = dduals.reshape(shape)
            if op=='add':
                self.duals += dduals
            elif op=='replace':
                self.duals = dduals
                
        if op=='add':
            self.dofs += ddofs
        elif op=='replace':
            self.dofs = ddofs
    
        return self.getDofs()
    
    def saveFields(self, key, fields, components):
        '''
        
        '''
        values = []
        args = components.__iter__()
        for f in fields:
            if not self.fields.hasField(f): 
                self.fields.addField(f)
            values.append(self.getCurrentField(f, args.next()))

        self.fields.addState(key, fields, values)
            
    def saveQOI(self, key, quantities, arguments):
        
        values = []
        args = arguments.__iter__()
        for q in quantities:
            if not self.quantities.hasQuantity(q): 
                self.quantities.addQuantity(q)
            value, complement = self.getCurrentQuantity(q,args.next())
            values.append(value)
        
        self.quantities.addState(key, quantities, values, complement)
        
        
    def computeOperators(self,*args):
        '''
        
        '''
        raise NotImplementedError
        
    def computeLhs(self,*args):
        '''
        
        '''
        raise NotImplementedError
    

        
    def computeRhs(self,*args):  
        '''
        
        '''  
        raise NotImplementedError
    
    def B(self,x):
        '''
        
        '''
        
        return B

    def printSavedFieldsAndQOI(self):
        
        print self.fields
        print self.quantities

## Particularized Meshes:
class DefaultMeshDecorator(DefaultMesh):
    '''
        Decorator of the DefaultMesh class. Enable to particularize a Mesh in a dynamic manner. 
    '''
    def __init__(self,defaultmesh):
        kwargs = copy.copy(defaultmesh.__dict__)
        UId=kwargs["UId"]
        del kwargs["UId"]
        super(DefaultMeshDecorator, self).__init__(UId,**kwargs)
          

class Mesh(DefaultMeshDecorator):
    '''
       Self-sufficient mesh. It is able to build operators from assembly of elementary contributions.
    '''
    
    def __init__(self, defaultmesh):
        '''
        Mesh constructor
        '''
        super(Mesh, self).__init__(defaultmesh)
        

    def computeLhs(self, LHS):
        
        n_enr = len(self.enrichments)+1
        n_nodes = len(self.nodesMap["C2U"])
        # initialize Lhs matrix
        
        lhs = sp.lil_matrix((6*n_nodes*n_enr,6*n_nodes*n_enr), dtype = FLOAT)
        
        for el in self.elementsMap["U2C"]:
            
            P = self.elementsFrame[el]
            Ae = self.elementsAO[el]
            
            if n_enr > 1:
                pass
            
            lhs_e = ELTYPE_TO_MODEL_CLASS[self.elementsType[el]].LHS_NAME_TO_METHOD[LHS](self.nodesCoo[self.elementsNodes[el],:],)
            lhs_e = np.dot(P.T,np.dot(lhs_e,P))
            
            lhs += np.dot(Ae,np.dot(lhs_e,Ae.T))
            
        self.lhs[LHS] = lhs
        
    
    def computeRhs(self, RHS):
        pass
    

# Mesh related functions:
    
    def consolidateMesh(self):
        ''' Check if all the mesh entities have the required properties defined. Calculates missing
        ones.
        '''
        # nodes:
        n_nodes = len(self.nodesMap["C2U"])
        self.nodesCoo = sp.csr_matrix(self.getNodesArray())
        
        # elements:
        n_elem = len(self.elementsMap["C2U"])  
        
        # build element frames:
        if self.elementsFrame == []:
            self.elementsFrame = [None] * n_elem
            for i,el in enumerate(self.elementsMap["C2U"]):
                self.elementsFrame[i] = ELTYPE_TO_MODEL_CLASS[self.elementsType[i]].frame(self.nodesCoo[self.elementsNodes[i],:])
        
            
        # set element properties:
        if self.elementsProp == []:
            self.elementsProp = [None] * n_elem
            for p,prop in enumerate(self.properties):
                for elset in prop['ELSET']:
                    for el in self.elemSet[self.elemSetMap["U2C"][elset]]:
                        self.elementsProp[el] = p    
        
        # initialize dofs
        if self.dofs == []:
            n_enr = len(self.enrichments)+1
            self.dofs = sp.lil_matrix((6*n_nodes*n_enr,1), dtype = FLOAT)

        # define elements 'elementary' assembly operators:
        if self.elementsAO == []:
            self.setElementsAO()


class InterfacedMesh(DefaultMeshDecorator):
    '''
       Mesh interfaced with third-party software. The computation of some operators is delegated to third-party software.
    '''
    

    def __init__(self, defaultmesh, parser=None, software='nastran' ): 
        
        available_softwares = {'nastran': Structural_Part,
                           }
        super(InterfacedMesh, self).__init__(defaultmesh)
        if parser:
            self.parser = parser
        elif software in available_softwares:
            self.parser = available_softwares[software](self.UId)
            self.parser.buildData(mesh=self)
        self.software = software
        self.dofsMap = None
        self.buildDofsMap()
        
    def __str__(self):
        
        r = super(InterfacedMesh, self).__str__()
        r += '\n' + '  - interface software: ' + str(self.software)
        return r 
        
    def computeOperators(self, lhslist = [], rhslist = [], disp=1, load=1):
        '''
            @param lhsList: list.
            @param rhsList: list.
        '''
        filenames={'dof':None,'input':'mesh-'+str(self.UId),'lhs':[],'rhs':[]}
        
        param = parameters.load()
        outdir = param.WORKDIR
        
        # write the dofs to a file
        if param.UPDATE_DOFS:
            filenames['dof']=outdir+'dof.'+str(self.UId)
            dofs={'DOF':None}
            dofs['DOF']=self.getDofs()
            # mapping between user's (i.e. third-party software) and icafe's numbering of d.o.f. 
            dofs['DOF']=self.dofsMap['U2C']*dofs['DOF']
            # write d.o.f. state to file
            self.parser.writeOperators(filenames['dof'],dofs)
            
        # write the bdf file: Assign variables, solution, DMAP alter, model data
        
        if len(lhslist)+len(rhslist)>0:
            for lhs in lhslist: 
                filenames['lhs'].append(outdir+lhs+str(self.UId))
            for rhs in rhslist:
                if rhs=='Fext': # External forces are managed by icafe
                    filenames['rhs'].append(None)
                else:
                    filenames['rhs'].append(outdir+rhs+str(self.UId))
            
        if param.UPDATE_INPUTFILE:
            # write the headers with DMAP options
            self.parser.buildHeaders(filenames,lhslist,rhslist)
            self.parser.buildData(mesh=self, updateMesh=False, updateDofs=param.UPDATE_DOFS)
            inputfile = self.parser.writeInputFile(filenames['input'],out=outdir)
        
        # run Nastran with bdf file.
        if param.RUN_THIRD_SOFTWARE:
            if param.DEBUG:
                print 'Mesh ', self.UId, ' calls 3rd-party software ', self.software,' with ', inputfile
            self.parser.runThirdSoftware(inputfile,outdir)
        
        # read the files containing operators.
        if not param.READ_OPERATORS:
            return {} 
        
        # read lhs operators
        for i,lhs in enumerate(lhslist):
#            print outdir
            op=self.parser.readOperators(filenames['lhs'][i],rename_operator=lhs)
#            print op[lhs].get_shape()
#            print self.dofsMap['U2C'].get_shape()
            self.lhs[lhs]=self.dofsMap['U2C']*op[lhs]*self.dofsMap['C2U']
        
        for i,rhs in enumerate(rhslist):
            if rhs=='Fext':
                try:
                    self.rhs[rhs]=self.constraints.getFext(load)
                except AttributeError: # no constraints
                    dofs = self.getDofs()
                    self.rhs[rhs]=sp.lil_matrix(dofs.get_shape())
            elif rhs=='R':
                op=self.parser.readOperators(filenames['rhs'][i],rename_operator=rhs)
                self.rhs['Fext']=self.constraints.getFext(load)
                self.rhs[rhs]=op[rhs] - self.rhs['Fext']
            else:
                op=self.parser.readOperators(filenames['rhs'][i],rename_operator=rhs)
                self.rhs[rhs]=self.dofsMap['U2C']*op[rhs]
        
        return self.lhs, self.rhs
    
    def computeRhs(self):
        
        pass
    
    def buildDofsMap(self):
        '''
            Building of the mapping between user's and icafe's numbering of d.o.f.
        
        '''
        # transformation matrix:
        A = sp.lil_matrix((self.ndofs_per_node*self.nnodes,self.ndofs_per_node*self.nnodes),dtype=FLOAT)
        
        sortedNodesMap = sorted(self.nodesMap['C2U'])
        
        for nid, nuid in enumerate(self.nodesMap['C2U']):
            
            nuid = sortedNodesMap.index(nuid)
            
            try:
                for dof in xrange(self.ndofs_per_node):
                    A[nid*self.ndofs_per_node + dof, nuid*self.ndofs_per_node + dof] = 1.0
                    
            except ValueError:
                print 'User id of node is not an integer: ', nuid  
                raise ValueError
        
        self.dofsMap = {'C2U':A, 'U2C':A.transpose()}     
#        self.ixDofsMap = np.ix_(self.dofsMap['C2U'],self.dofsMap['C2U'])

#class NastranMesh(InterfacedMesh):
#    
#    def __init__(self,):


class Constraints():
    '''
        An object containing kinematic and static constraints on a mesh. 
    '''
    def __init__(self, Cd, d, Cf, f, meshId=None):
        '''
            Constructor.
            @param Cd: List of sparse matrices. Extractors of degrees of freedom, one per load case.
            @param d: List of sparse vectors. Values of conditions on d.o.f. for each load case.
            @param Cf: List of sparse matrices. Extractors of forces, one per load case.
            @param f: List of sparse vectors. Values of forces for each load case.
            @param meshId: str, optional. Cross reference to mesh.
        '''
        self.Cd = Cd
        self.d = d
        self.Cf = Cf
        self.f = f
        self.meshId = meshId
        
    def __str__(self):
        
        r = '   - Constraint matrices:\n'
        Cd, d, Cf, f = self.get()
        r += str(Cd[0].data) + '\n'
        r += str(d[0].data) + '\n'
        r += str(Cf[0].data) + '\n'
        r += str(f[0].data)
        return r
        
    def get(self, loadcase=None):
        '''
            Returns contraints corresponding to loadcase.
            @param loadcase: int or None, optional. 
        
        '''
        if loadcase:
            return self.Cd[loadcase], self.d[loadcase], self.Cf[loadcase], self.f[loadcase]
        else:
            return self.Cd, self.d, self.Cf, self.f
        
    def getFext(self, load=1, loadcase=None):
        
        if loadcase:
            F = self.Cf[loadcase].T*self.f[loadcase]

        else:
            F = sp.bmat([[self.Cf[i].T*self.f[i] for i in xrange(len(self.Cf))]])
                   
        return load*F


# container classes of saved data: fields, quantities of interest, ...         
class Fields(object):
    
    def __init__(self, meshId, pickling=True):
        self.fields = {}
        self.statesMap = {}
        self.nStates = 0
        self.nFields = 0
        # Pickling Info
        self.pickled = False
        self.pickling = pickling  # may be set by user in parameter file
        self.meshId = meshId
        param = parameters.load()
        
        self.filename = param.WORKDIR + 'fields-mesh-'+str(self.meshId)+'.pkl'
    
    def __str__(self):
        
        msg = str(self.nFields) + ' Fields in ' + str(self.nStates) + ' States saved in mesh ' + str(self.meshId) + ':\n'
        
        return msg
    
    def getField(self,field,state=None):
        self.unPickle()
        if not state:
            try:
                self.pickle()
                return self.fields[field]
            except:
                self.pickle()
                return None
        else:
            field = self.fields[field]
            value = field.getState(self.statesMap[state])
            self.pickle()
            return value
            
    def hasField(self, field):
        self.unPickle()
        msg = self.fields.has_key(field)
        self.pickle()
        return msg
        
    def addField(self, field):
        self.unPickle()
        self.fields[field] = Field(field, self.nStates)
        self.nFields += 1
        self.pickle()
    
    def rmField(self, field):
        self.unPickle()
        del self.fields[field]
        self.nFields -= 1
        self.pickle()
    
    def addState(self, key, fields, values):
        
        self.unPickle()
        
        if not self.statesMap.has_key(key):
            self.statesMap[key] = copy.copy(self.nStates)
            index=None
            self.nStates += 1
        else:
            index = self.statesMap[key]
            
        for f in self.fields.keys():
            try:
                self.fields[f].addState(values[fields.index(f)], index)
            except ValueError:
                self.fields[f].addState(None)
            except IndexError:
                raise IndexError('Bad values given to Fields.addState(key, fields, values): values = ' + str(values))
        
        
        
        self.pickle()
    
    def pickle(self):
        
        f = open(self.filename,'wb')
        pickler.dump(self.__dict__,f)
        f.close()
        self.fields = 'PICKLED'
        self.statesMap = 'PICKLED'
        self.pickled = True
        
    def unPickle(self):
        
        if self.pickled:
            data = pickler.load(file(self.filename,'rb'))
            self.fields = data.get('fields',{})
            self.statesMap = data.get('statesMap',{})
    
    
class Field(object):
    '''
       Storage of the states of a field.
    '''
    def __init__(self, name, nstates=0):
        if nstates > 0:
            self.states = [None]*nstates
        else:
            self.states = []
        
        self.name = name
    
    def addState(self, val, index=None):
        '''
        
        '''
        if not index:
            self.states.append(copy.copy(val))
        else:
            self.states[index]=val
    
    def getState(self, index):
        
        return self.states[index]        

class Quantities(object):
    
    def __init__(self, meshId, pickling=True):
        self.quantities = {}
        self.statesMap = {}
        self.nStates = 0
        self.nQoi = 0
        # Pickling Info
        self.pickled = False
        self.pickling = pickling  # may be set by user in parameter file
        self.meshId = meshId
        param = parameters.load()
        
        self.filename = param.WORKDIR + 'qoi-mesh-'+str(self.meshId)+'.pkl'
    
    def __str__(self):
        
        self.unPickle()
        param = parameters.load()
        
        msg = str(self.nQoi) + ' QOI in ' + str(self.nStates) + ' States saved in mesh ' + str(self.meshId) + ':\n'
        msg+= 'states, ' + str(self.statesMap.keys()).strip('[').strip(']') + '\n'
        for qoi in self.quantities.values():
            
            if not qoi.active: continue
            
            for c in range(qoi.components):
                if qoi.components>1:
                    msg+= str(qoi.name)+'-'+str(c)
                else:
                    msg+= str(qoi.name)
                    
                for val in qoi.states:              # for each state
                    try:
                        msg+= ',' + str(val[c])
                    except:
                        msg += ',' + str(val)  
                msg += '\n'
                
        
        self.pickle()
                
        return msg
    
    def hasQuantity(self, qoi):
        self.unPickle()
        msg = self.quantities.has_key(qoi)
        self.pickle()
        return msg
        
    def addQuantity(self, qoi, replace=True):
        self.unPickle()
        if self.quantities.has_key(qoi) and replace:
            self.quantities[qoi] = Quantity(qoi, self.nStates)
            index = None
        elif not self.quantities.has_key(qoi):
            self.quantities[qoi] = Quantity(qoi, self.nStates)
            self.nQoi += 1
            index = -1
        else:
            index=None
        self.pickle()
        return index
    
    def rmQuantity(self, qoi):
        self.unPickle()
        del self.quantities[qoi]
        self.nQoi -= 1
        self.pickle()
    
    def getQuantities(self):
        self.unPickle()
        quantities = copy.copy(self.quantities)
        self.pickle()
        return quantities
    
    def addState(self, key, qois, values, complement=None):
        self.unPickle()
        
        if not self.statesMap.has_key(key):
            self.statesMap[key] = copy.copy(self.nStates)
            index=None
            self.nStates += 1
        else:
            index = self.statesMap[key]
        
        
        for q in self.quantities.keys():
            
            if '-' in q:
                continue
            
            try:
                if complement[qois.index(q)]:
                    
                    for j,k in enumerate(complement):
                        
                        qoi = q+'-'+str(k)

                        idx = self.addQuantity(qoi,replace=False)
                        if not idx:
                            idx = index
                        self.unPickle()

                        try:
                            self.quantities[qoi].addState(values[qois.index(q)][j], idx)
                        except ValueError:
                            self.quantities[qoi].addState(None)
                        except IndexError:
                            raise IndexError('Bad values given to quantities.addState(key, fields, values): values = ' + str(values))
                        
                            
    #                self.rmQuantity(q)
    #                self.unPickle()    
                    self.quantities[q].addState('unactivated', index)
                else:
                    self.quantities[q].addState(values[qois.index(q)], index)
            except ValueError:
                self.quantities[q].addState(None, index)
            except IndexError:
                raise IndexError('Bad values given to quantities.addState(key, fields, values): values = ' + str(values))
        
        self.pickle()
    
    def pickle(self):
        
        if not self.pickled:
            f = open(self.filename,'wb')
            pickler.dump(self.__dict__,f)
            f.close()
            self.quantities = 'PICKLED'
            self.statesMap = 'PICKLED'
            self.pickled = True
        
    def unPickle(self):
        
        if self.pickled:
            data = pickler.load(file(self.filename,'rb'))
            self.quantities = data.get('quantities',{})
            self.statesMap = data.get('statesMap',{})  
            self.pickled = False
                  
class Quantity(object):
    '''
       Storage of the states of a quantity of interest.
    '''
    def __init__(self, name, nstates=0):
        if nstates > 0:
            self.states = [None]*nstates
        else:
            self.states = []
        
        self.name = name
        self.components = 1
        self.active = True
        
    
    def addState(self, val, index=None):
        '''
        
        '''
        if val=='unactivated' or not self.active:
            self.active = False
            return 0
            
        if isinstance(val,list) and self.components==1:
            self.components = len(val)
        
        if not index:
            self.states.append(copy.copy(val))
        else:
            self.states[index]=val